import numpy as np

# 辅助函数：为每个数据点分配到最近的簇
def assign_cluster(data, centroids):
    # 计算每个数据点与每个质心的距离
    distances = np.linalg.norm(data[:, np.newaxis] - centroids, axis=2)
    # 为每个点分配最近的质心
    return np.argmin(distances, axis=1)

# KMeans算法实现
def kmeans(data, k, epsilon=1e-4, max_iterations=300):
    # 随机初始化k个质心
    centroids = data[np.random.choice(data.shape[0], k, replace=False)]
    
    for i in range(max_iterations):
        # 步骤1：分配每个点到最近的簇
        clusters = assign_cluster(data, centroids)
        
        # 步骤2：更新每个簇的质心
        new_centroids = np.array([data[clusters == i].mean(axis=0) for i in range(k)])
        
        # 步骤3：检查收敛条件（质心的变化小于epsilon）
        eps = np.linalg.norm(new_centroids - centroids)
        
        if eps < epsilon:
            print(f"Converged after {i + 1} iterations.")
            break
        
        centroids = new_centroids
    
    return centroids, clusters

# 示例用法
if __name__ == "__main__":
    # 生成一些示例数据
    data = np.random.rand(100, 2)  # 100个二维数据点
    
    # 设置KMeans参数
    k = 3
    epsilon = 1e-4
    max_iterations = 300
    
    centroids, clusters = kmeans(data, k, epsilon, max_iterations)
    
    print("最终的质心：", centroids)
    print("每个点的簇分配：", clusters)
